
% This routine is used to calculate the vertical wavenumber 
% spectra from MMP observations of:
%
% a) Kinetic energy
% b) Shear
% 
%
% Author: Ritabrata Thakur, 10 September 2021
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Mode of usage: fully automatic
%
% Last used: 26/1/22
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

MMP_no_array = {1,3,4};

for jj=1:length(MMP_no_array)  
% This automatically chooses the MMP number from the above array
MMP_number = MMP_no_array{jj}; 

clearvars -except jj MMP_no_array MMP_number 
fprintf('Running MMP_number %d\n',MMP_number);

%% Loading data
% this is depth and time range where the velocity is good (not NaN)
% for MMP1
if (MMP_number == 1)
 MMP_load=load('iwap06_MP1_WHOI112.mat');
 depth_start = 43;
 depth_end   = 700;
 lat = 25.5;   %North
 lon = 195;  %East
 start_time = 4;
 end_time = 980;
 inertial_start = 25.05; % 90% of 27.84 hrs
 inertial_end   = 30.6;  % 110% of 27.84 hrs
end
% for MMP3
if (MMP_number == 3)
 MMP_load=load('iwap06_MP3_WHOI114.mat');
 depth_start = 44;
 depth_end   = 702;
 lat = 28.9;   %North
 lon = 196.5;  %East
 start_time = 4;
 end_time = 980;
 inertial_start = 22.34; % 90% of 24.83 hrs
 inertial_end   = 27.31; % 110% of 24.83 hrs
end
% for MMP4
if (MMP_number == 4)
 MMP_load=load('iwap06_MP4_101.mat');
 depth_start = 44;
 depth_end   = 700;
 lat = 30.1;   %North
 lon = 197.1;  %East
 start_time = 4;
 end_time = 538;
 inertial_start = 21.52; % 90% of 23.92 hrs
 inertial_end   = 26.31; % 110% of 23.92 hrs
end

% reading the raw variables
mmp.U    = MMP_load.mp.u;
mmp.V    = MMP_load.mp.v;
mmp.temp = MMP_load.mp.t;
mmp.salt = MMP_load.mp.s;
mmp.depth= MMP_load.mp.z; 
mmp.time_array  = squeeze(MMP_load.mp.datenum1h);

%% Brunt Vaisala
mmp.Z = -mmp.depth(depth_start:depth_end);
mmp.p = gsw_p_from_z(squeeze(mmp.Z),lat);

[SA, ~] = gsw_SA_from_SP(mmp.salt(depth_start:depth_end,:),mmp.p,lon,lat);  %converts practical salinity
%to absolute salinity which is used by the N2 routine. The flag of "1" if
%the measurement is in ocean. If it is used in land, set it to 0. do not
%put anything. it is automatically set
CT = gsw_CT_from_t(SA,mmp.temp(depth_start:depth_end,:),mmp.p); %converts in-situ temperature to
%conservative temperature as is required in GSW toolbox

clear N2 NN
% Brunt Vaisala frequency
for i=1:length(CT(1,:))
 [N2,p_mid] = gsw_Nsquared(SA(:,i),CT(:,i),mmp.p,lat);
 NN(i,:)=N2;  %this is the total Brunt-vaisala frequency squared stored
 % at a single location. It is probably best to continue with one location
 % as averaging over different locations would mean we would have to
 % calculate the wkb factor for each depth differently as the
 % stratification would vary according to location
end

% removing the values where there is unstable stratification
for i=1:length(NN(:,1))
    for j=1:length(NN(1,:))
      if(NN(i,j) <= 0)
      NN(i,j) = NaN;
      end
    end
end

%sqrt of the N^2
BV = sqrt(NN);

%interpolating to original depth levels. As the original units are in
%pressure (p), so interpolating it to that level would amount to
%interpolating to original depth
% for i=1:length(NN(:,1))
%  BV_interp(i,:) = interp1(p_mid, (BV(i,:)), mmp.p, 'pchip');
% end

for i=1:length(NN(:,1))
 BV_interp(i,:) = interp1(p_mid, BV(i,:), mmp.p, 'linear', 'extrap');
end

% using the depth range where the data is good
mmp.depth_array_short = squeeze(mmp.depth(depth_start:depth_end));
BV_interp_short   = squeeze(BV_interp(:,:));
BV_interp_time_mean =  mean(BV_interp_short, 'omitnan');
BV_interp_time_mean = naninterp(BV_interp_time_mean);

N_0_mean = 0.0052; %used in literature = 3cph

mmp.wkb_factor_short = sqrt(BV_interp_time_mean./N_0_mean);

%% WKB scaling

mmp.z_stretched(1) = mmp.depth_array_short(1);
for ii=2:length(mmp.depth_array_short)
 mmp.z_stretched(ii) = mmp.z_stretched(1) + trapz(mmp.depth_array_short(1:ii), squeeze(BV_interp_time_mean(1:ii)))./N_0_mean; 
end

% this array is needed to interpolate the velocities from unequally-spaced
% stretched coordinates to equally-spaced
mmp.z_stretched_equal_spaced = linspace(mmp.z_stretched(1), mmp.z_stretched(end), length(mmp.z_stretched(1:end)));

% WKB like scaling of the horizontal velocities at the intermediate depth
% where the stratification does not vary abruptly
for k=1:length(mmp.U(1,:))
    mmp.wkb_scaled.U(:,k) = squeeze(mmp.U(depth_start:depth_end,k))'./mmp.wkb_factor_short;
    mmp.wkb_scaled.V(:,k) = squeeze(mmp.V(depth_start:depth_end,k))'./mmp.wkb_factor_short;
end    

% We now have the stretched fields defined in "z_stretched" points but we
% want to interpolate to "z_stretched_equal_spaced" for ease of psd
% calculation. I am in favour of calculating spectra in equal spacing
% rather than Lomb-Scargle because of added complexity of the latter and
% also because in the original Leaman and Sanford, they do what I have done
%, i.e. interpolate to equally spaced grid before psd calculation
[grid_old_z, grid_old_time] = ndgrid(squeeze(mmp.z_stretched), mmp.time_array);
[grid_new_z, grid_new_time] = ndgrid(squeeze(mmp.z_stretched_equal_spaced), mmp.time_array);
%Only doing a 2D interpolation here as the data has only two directions now
mmp.wkb_scaled_equal_z.U = interpn(grid_old_z, grid_old_time, mmp.wkb_scaled.U, grid_new_z, grid_new_time);
mmp.wkb_scaled_equal_z.V = interpn(grid_old_z, grid_old_time, mmp.wkb_scaled.V, grid_new_z, grid_new_time);
%W_scaled_equal_z = interpn(grid_old_x, grid_old_y, grid_old_z, grid_old_time, W, grid_new_x, grid_new_y, grid_new_z, grid_new_time);

%removing any NaNs
for i=1:length(mmp.wkb_scaled_equal_z.U(:,1))
  mmp.wkb_scaled_equal_z.U(i,:) = naninterp(mmp.wkb_scaled_equal_z.U(i,:));
  mmp.wkb_scaled_equal_z.V(i,:) = naninterp(mmp.wkb_scaled_equal_z.V(i,:));
end

N = length(mmp.z_stretched_equal_spaced);
win = hann(N);

% Number of FFT points for the Fourier transform
NFFT = N; 

%Fs=2*pi/(mmp.z_stretched_equal_spaced(2)-mmp.z_stretched_equal_spaced(1)); %removed 2pi on 04Nov21
%following early literature (Gregg et al 1991: "varieties of fully resolved spectra of vertical shear")
% this 2pi is also removed from shear and strain
Fs = 1/(mmp.z_stretched_equal_spaced(2)-mmp.z_stretched_equal_spaced(1)); 
mmp.waveno = 0:Fs/NFFT:Fs/2; % we can do this because the depth
% array is equally spaced albeit in stretched coordinates

% getting the time-averaged psd from the following. I am also getting
% "psdx_array" which is dependent on time too. This has been output to
% understand the time evolution of the psd

%% Filtering velocities in hourly bands and near-inertial
% bands to calculate shear in those hourly bands. 
for i=1:length(mmp.wkb_scaled_equal_z.U(:,1))
  %mmp.wkb_scaled_equal_z.U_3_to_5hr(i,:)  = bandpass(squeeze(mmp.wkb_scaled_equal_z.U(i,start_time:end_time)),[1/5.20 1/2.80], 1); 
  mmp.wkb_scaled_equal_z.U_6_to_8hr(i,:)  = bandpass(squeeze(mmp.wkb_scaled_equal_z.U(i,start_time:end_time)),[1/8.20 1/5.80], 1); 
  mmp.wkb_scaled_equal_z.U_sem_diur(i,:)  = bandpass(squeeze(mmp.wkb_scaled_equal_z.U(i,start_time:end_time)),[1/13.5 1/11.5], 1); 
  mmp.wkb_scaled_equal_z.U_lowpass(i,:)   = lowpass(squeeze(mmp.wkb_scaled_equal_z.U(i,start_time:end_time)),1/11.5, 1);
  mmp.wkb_scaled_equal_z.U_highpass(i,:)  = highpass(squeeze(mmp.wkb_scaled_equal_z.U(i,start_time:end_time)),1/11.5, 1);   
  mmp.wkb_scaled_equal_z.U_inertial(i,:)  = bandpass(squeeze(mmp.wkb_scaled_equal_z.U(i,start_time:end_time)),[1/inertial_end 1/inertial_start], 1); 
  %mmp.wkb_scaled_equal_z.V_3_to_5hr(i,:)  = bandpass(squeeze(mmp.wkb_scaled_equal_z.V(i,start_time:end_time)),[1/5.20 1/2.80], 1); 
  mmp.wkb_scaled_equal_z.V_6_to_8hr(i,:)  = bandpass(squeeze(mmp.wkb_scaled_equal_z.V(i,start_time:end_time)),[1/8.20 1/5.80], 1); 
  mmp.wkb_scaled_equal_z.V_sem_diur(i,:)  = bandpass(squeeze(mmp.wkb_scaled_equal_z.V(i,start_time:end_time)),[1/13.5 1/11.5], 1);
  mmp.wkb_scaled_equal_z.V_lowpass(i,:)   = lowpass(squeeze(mmp.wkb_scaled_equal_z.V(i,start_time:end_time)),1/11.5, 1);
  mmp.wkb_scaled_equal_z.V_highpass(i,:)  = highpass(squeeze(mmp.wkb_scaled_equal_z.V(i,start_time:end_time)),1/11.5, 1);   
  mmp.wkb_scaled_equal_z.V_inertial(i,:)  = bandpass(squeeze(mmp.wkb_scaled_equal_z.V(i,start_time:end_time)),[1/inertial_end 1/inertial_start], 1); 
end


%% KE psd. Total PSD in different time bands
clear mmp.psd_time_avg mmp.psdx_array
% taking psd from time location where there is no NaN. The time step from 4
% to 980 does not have NaNs and hence produces real psd. Otherwise it would
% produce complex psds
%
% 26 Jan 2022: Calculating KE psd at different time bands because it would be
% needed to calculate shear using Gregg estimate
[mmp.psd_time_avg(:), mmp.psdx_array(:,:)] = welch_estimate_corrected_rt((mmp.wkb_scaled_equal_z.U(:,start_time:end_time))...
    ,(mmp.wkb_scaled_equal_z.V(:,start_time:end_time)),win,NFFT,Fs); 

[mmp.psd_6_8hr(:), ~]    = welch_estimate_corrected_rt((mmp.wkb_scaled_equal_z.U_6_to_8hr(:,:))...
    ,(mmp.wkb_scaled_equal_z.V_6_to_8hr(:,:)),win,NFFT,Fs);

[mmp.psd_sem_diur(:), ~] = welch_estimate_corrected_rt((mmp.wkb_scaled_equal_z.U_sem_diur(:,:))...
    ,(mmp.wkb_scaled_equal_z.V_sem_diur(:,:)),win,NFFT,Fs);

[mmp.psd_lowpass(:), ~]  = welch_estimate_corrected_rt((mmp.wkb_scaled_equal_z.U_lowpass(:,:))...
    ,(mmp.wkb_scaled_equal_z.V_lowpass(:,:)),win,NFFT,Fs);

[mmp.psd_highpass(:), ~] = welch_estimate_corrected_rt((mmp.wkb_scaled_equal_z.U_highpass(:,:))...
    ,(mmp.wkb_scaled_equal_z.V_highpass(:,:)),win,NFFT,Fs);

[mmp.psd_inertial(:), ~] = welch_estimate_corrected_rt((mmp.wkb_scaled_equal_z.U_inertial(:,:))...
    ,(mmp.wkb_scaled_equal_z.V_inertial(:,:)),win,NFFT,Fs);

% there is no lat-lon average here as the Mclane profiler is only for a
% specific location and does not have any spatial variability
mmp.psd_time_avg  = mmp.psd_time_avg./(NFFT*Fs);
mmp.psd_6_8hr     = mmp.psd_6_8hr./(NFFT*Fs);
mmp.psd_sem_diur  = mmp.psd_sem_diur./(NFFT*Fs);
mmp.psd_lowpass   = mmp.psd_lowpass./(NFFT*Fs);
mmp.psd_highpass  = mmp.psd_highpass./(NFFT*Fs);
mmp.psd_inertial  = mmp.psd_inertial./(NFFT*Fs);
mmp.comment       = 'The MMP depth range is from 83-1348m for the whole period from 25 April 2006';
mmp.psdx_array_normalised    = mmp.psdx_array./(NFFT*Fs);

%% Bootstrap intervals on means
% each row of bootstat contains the mean and std of a bootstrap sample
% psd_ci is a 3D matrix that contains the upper and lower bound of 95%
% (default) confidence intervals at each vertical wavenumber (the third
% dimension)
for j = 1:length(mmp.psdx_array_normalised(1,:))
 [psd_ci(:,:,j),bootstat] = bootci(100,@(x)[mean(x) std(x)],mmp.psdx_array_normalised(:,j)); %,'alpha',0.05,'type','norm'
end

mmp_psd_ke.psd_ci = psd_ci;
mmp_psd_ke.bootstat = bootstat;

%% Raw standard deviation to get another estimate of error bars. This is not
% the same as the one above as that calculates the 95% confidence intervals
% whereas this is just raw standard deviation at each vertical wavenumber.

for j = 1:length(mmp.psdx_array_normalised(1,:))
 stdev(j) = std(mmp.psdx_array_normalised(:,j));
end

mmp_psd_ke.stdev = stdev;

%% saving the KE psds
mmp_psd_ke.ke_waveno  = mmp.waveno; 
mmp_psd_ke.ke_tot_psd = mmp.psd_time_avg;
mmp_psd_ke.ke_6_8hr   = mmp.psd_6_8hr;
mmp_psd_ke.ke_sem_diur= mmp.psd_sem_diur;
mmp_psd_ke.ke_lowpass = mmp.psd_lowpass;
mmp_psd_ke.ke_highpass= mmp.psd_highpass;
mmp_psd_ke.ke_inertial= mmp.psd_inertial; 
mmp_psd_ke.psd_array  = mmp.psdx_array_normalised;
mmp_psd_ke.comment    = mmp.comment;
savename              = [basedir 'mmp' sprintf('%d', MMP_number) '/psd_ke.mat'];
save(savename, 'mmp_psd_ke');

%% Shear psd Pinkel estimate
% update on 19 Jan 2022 to save the shear spectra along with the ke spectra

clear psdx_array_shear mmp.shear_psd_avg
%%%%%%%%%%%%%%%%
% calling the waveno-freq calculation routine. Complex velocity is calculated
% inside the following routine. This is a single-sided spectra and hence
% has a factor 2 multiplied inside routine. The routine is based on
% Pinkel's shear estimate
[psdx_array_shear,mmp.pinkel] = waveno_only_spectra_shear_ybp(mmp.wkb_scaled_equal_z.U(:,start_time:end_time),...
    mmp.wkb_scaled_equal_z.V(:,start_time:end_time),squeeze(mmp.z_stretched_equal_spaced),NFFT);
%%%%%%%%%%%%%%%%

waveno_shear                     = 0:Fs/(NFFT-1):Fs/2;  
mmp_psd_shear.pinkelwaveno       = waveno_shear;
% 1 is subtracted from NFFT because the FFTs are taken in that many number
% of points in case of Pinkel estimate
mmp_psd_shear.pinkel_tot_shear   = mmp.pinkel./((NFFT-1)*Fs);

% estimates of shear in different time bands. Added on 25 Jan 2022. These
% velocities are already in the start_time:end_time period and hence we
% can't specify that here.
[~,mmp.pinkel_shear_6_8hr]       = waveno_only_spectra_shear_ybp(mmp.wkb_scaled_equal_z.U_6_to_8hr,...
    mmp.wkb_scaled_equal_z.V_6_to_8hr,squeeze(mmp.z_stretched_equal_spaced),NFFT);
[~,mmp.pinkel_shear_sem_diur]    = waveno_only_spectra_shear_ybp(mmp.wkb_scaled_equal_z.U_sem_diur,...
    mmp.wkb_scaled_equal_z.V_sem_diur,squeeze(mmp.z_stretched_equal_spaced),NFFT);
[~,mmp.pinkel_shear_highpass]    = waveno_only_spectra_shear_ybp(mmp.wkb_scaled_equal_z.U_highpass,...
    mmp.wkb_scaled_equal_z.V_highpass,squeeze(mmp.z_stretched_equal_spaced),NFFT);
[~,mmp.pinkel_shear_inertial]    = waveno_only_spectra_shear_ybp(mmp.wkb_scaled_equal_z.U_inertial,...
    mmp.wkb_scaled_equal_z.V_inertial,squeeze(mmp.z_stretched_equal_spaced),NFFT);

mmp_psd_shear.pinkel_shear_6_8hr    = mmp.pinkel_shear_6_8hr./((NFFT-1)*Fs);
mmp_psd_shear.pinkel_shear_sem_diur = mmp.pinkel_shear_sem_diur./((NFFT-1)*Fs);
mmp_psd_shear.pinkel_shear_highpass = mmp.pinkel_shear_highpass./((NFFT-1)*Fs);    
mmp_psd_shear.pinkel_shear_inertial = mmp.pinkel_shear_inertial./((NFFT-1)*Fs);


%% Gregg method of calculating shear. This calculation of shear spectra 
% from velocity spectra is based on Gregg, Winkel and Sanford 1992.
% As this estimate is based on kinetic energy spectra, it has the same
% wavenumbers as the kinetic energy

 mmp_psd_shear.gregg_waveno         =  mmp.waveno;
 mmp_psd_shear.gregg_tot_shear      = (2*pi.*mmp_psd_shear.gregg_waveno).^2.*mmp_psd_ke.ke_tot_psd; 
 mmp_psd_shear.gregg_shear_6_8hr    = (2*pi.*mmp_psd_shear.gregg_waveno).^2.*mmp_psd_ke.ke_6_8hr; 
 mmp_psd_shear.gregg_shear_sem_diur = (2*pi.*mmp_psd_shear.gregg_waveno).^2.*mmp_psd_ke.ke_sem_diur; 
 mmp_psd_shear.gregg_shear_lowpass  = (2*pi.*mmp_psd_shear.gregg_waveno).^2.*mmp_psd_ke.ke_lowpass;
 mmp_psd_shear.gregg_shear_highpass = (2*pi.*mmp_psd_shear.gregg_waveno).^2.*mmp_psd_ke.ke_highpass; 
 mmp_psd_shear.gregg_shear_inertial = (2*pi.*mmp_psd_shear.gregg_waveno).^2.*mmp_psd_ke.ke_inertial; 
 
 %% saving the shear psds - both Pinkel and Gregg estimates with total
 % and band-averaged estimates
basedir_shear = /home;
savename = [basedir_shear 'mmp' sprintf('%d', MMP_number) '/psd_shear.mat'];
if exist(savename, 'file')
    delete(savename);
end
save(savename, 'mmp_psd_shear');
end
